package ai.george.rag.controller;


import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;

import java.util.List;
import java.util.UUID;

@Tag(name = "Qqrant向量数据库")
@RestController
@RequestMapping("/vector")
public class VectorController {

    @Autowired
    private VectorStore vectorStore;





    @Operation(summary = "向量数据库添加数据")
    @PostMapping("/add")
    public String add(@RequestParam(name = "text") String text) {
        try {
            Document document = Document.builder()
                    .id(UUID.randomUUID().toString()) // 可以手动设置id，不设置则自动生成一个随机字符串作为id
                    .text(text)
                    .build();
            vectorStore.add(List.of(document));
//            System.out.println(document);
            return "success";
        } catch (Exception e) {
            throw e;
        }
    }


    @Operation(summary = "向量数据库查询数据")
    @GetMapping("/query")
    public List<Document> query(@RequestParam(name = "query") String query) {
        try {
            SearchRequest request = SearchRequest.builder()
                    .query(query)
                    .topK(3)
                    .build();

            List<Document> documents = vectorStore.similaritySearch(request);
            return documents;
        } catch (Exception e) {
            throw e;
        }
    }
}
